#!/usr/local/bin/python3
# coding=utf-8

# <bitbar.title>Wuhan pneumonia data</bitbar.title>
# <bitbar.version>v0.11</bitbar.version>
# <bitbar.author>Yifen Wu</bitbar.author>
# <bitbar.author.github>Anthonyeef</bitbar.author.github>
# <bitbar.desc>Wuhan pneumonia is spreading in the world, mainly in China. This plugin will show information (people having pneumonia, people dead because of pneumonia, and people who are cured from this pneumonia) for each province in China.</bitbar.desc>
# <bitbar.image>https://tva1.sinaimg.cn/large/006tNbRwly1gbccqabcaoj30tw0lc4l6.jpg</bitbar.image>
# <bitbar.dependencies>python</bitbar.dependencies>
# <bitbar.abouturl>https://github.com/Anthonyeef/wuhan-virus-bitbar-plugin</bitbar.abouturl>

# 武汉加油

import requests
import re
import json
import os
import time
import sys

# 填写想看到的省份的名字，如
# targetProvinceName = {"北京", "湖北", "广东"}
targetProvinceName = {}

# 如果targetProvinceName不填，默认展示确诊人数前多少的省份
frontCount = 3

# 除了 targetProvinceName 之外，还想额外看到的省份
# 如果不填则不会展示
additionProvinceName = {"河南", "广西", "浙江", "江西"}

# 上次数据是否备份
backup = True

# 变化趋势颜色
increasedColor = '#FF0000'
noChangeColor = '#FFFFFF'
decreasedColor = '#00BFFF'

# 是否下一级菜单显示变化数据
countryChangeNextMenu = True
provinceChangeNextMenu = True
cityChangeNextMenu = False

confirmedString = '确诊'
suspectedString = '疑似'
deadString = '死亡'
curedString = '治愈'

altsep = os.path.sep
linesep = os.linesep
dataRootPath = '%s%sdata%s' % (sys.path[0], altsep, altsep,)
dataPath = '%s2019_nCoV%s' % (dataRootPath, altsep)
lastFile = '%slast.json' % dataPath

jsondata = json.loads('[]')
displayString = ''
updateString = ''
lastDataJSON = json.loads('{}')

bitBarDarkMode = os.getenv('BitBarDarkMode', 0)
textColor = "black"
if bitBarDarkMode:
    textColor = "white"


def initLastData():
    result = 0
    if not os.path.exists(dataRootPath):
        os.mkdir(dataRootPath)
        result = 1
    if not os.path.exists(dataPath):
        os.mkdir(dataPath)
        result = 1
    if not os.path.exists(lastFile):
        with open(lastFile, 'w') as file:
            file.write('{}')
        result = 1
    return result


def readLastData():
    global lastDataJSON
    with open(lastFile, 'r') as file:
        line = file.readline()
        lastDataJSON = json.loads(line)


def getRecordFormatTime():
    result = ''
    try:
        result = time.strftime("%Y年%m月%d日%H时%M分%S秒", time.localtime(lastDataJSON.get('record_time')))
    except Exception:
        pass
    return result


def writeLastData():
    global updateString

    if initLastData() == 0 & backup:
        with open('%s%s.json' % (dataPath, getRecordFormatTime()), 'w') as file:
            file.write(str(lastDataJSON).replace("'", '"'))
    data = '{"record_time":%d%s}' % (time.time(), updateString)
    with open(lastFile, 'w') as file:
        file.write(data)


def getData(countryName, cityName, provinceName, dataName):
    result = 0
    try:
        if countryName is not None and countryName != '':
            if cityName is None or cityName == '':
                result = lastDataJSON.get(countryName).get(dataName)
            elif provinceName is None or provinceName == '':
                result = lastDataJSON.get(countryName).get(cityName).get(dataName)
            else:
                result = lastDataJSON.get(countryName).get(cityName).get(provinceName).get(dataName)
    except:
        result = None
    if result is None:
        result = 0
    return result


def getDifferences(countryName, cityName, provinceName, confirmedCount, suspectedCount, deadCount, curedCount):
    lastConfirmedCount = getData(countryName, cityName, provinceName, confirmedString)
    lastSuspectedCount = getData(countryName, cityName, provinceName, suspectedString)
    lastDeadCount = getData(countryName, cityName, provinceName, deadString)
    lastCuredCount = getData(countryName, cityName, provinceName, curedString)

    resultConfirmedCount = (confirmedCount - lastConfirmedCount if (lastConfirmedCount != 0) else 0)
    resultSuspectedCount = (suspectedCount - lastSuspectedCount if (lastSuspectedCount != 0) else 0)
    resultDeadCount = (deadCount - lastDeadCount if (lastDeadCount != 0) else 0)
    resultCuredCount = (curedCount - lastCuredCount if (lastCuredCount != 0) else 0)
    return {confirmedString: resultConfirmedCount, suspectedString: resultSuspectedCount, deadString: resultDeadCount, curedString: resultCuredCount}


def getColor(count):
    global increasedColor
    global noChangeColor
    global decreasedColor

    result = ' | color='
    if count > 0:
        result += increasedColor
    elif count == 0:
        result += noChangeColor
    else:
        result += decreasedColor
    return result


def getSplit(menuCount, start):
    result = start
    for i in range(menuCount):
        result += '--'
    return result


def getMenuSeq(menuCount):
    return getSplit(menuCount, '')


def getPartLine(menuCount):
    return getSplit(menuCount, '---')


def getSignedNumber(number):
    result = abs(number)
    if number > 0:
        result = '+%d' % result
    elif number < 0:
        result = '-%d' % result
    return result


def getgetDifferencesString(countryName, cityName, provinceName, confirmedCount, suspectedCount, deadCount, curedCount,
                            menuCount, nextMenu):
    global confirmedString
    global suspectedString
    global deadString
    global curedString

    result = ''
    if nextMenu:
        result += '%s变化%s' % (getMenuSeq(menuCount), linesep)
        menuCount += 1
    differences = getDifferences(countryName, cityName, provinceName, confirmedCount, suspectedCount, deadCount, curedCount)
    result += ('%s%s: %s%s%s' % (
        getMenuSeq(menuCount), confirmedString, getSignedNumber(differences[confirmedString]), getColor(differences[confirmedString]), linesep))
    result += ('%s%s: %s%s%s' % (
        getMenuSeq(menuCount), suspectedString, getSignedNumber(differences[suspectedString]), getColor(differences[suspectedString]), linesep))
    result += ('%s%s: %s%s%s' % (
        getMenuSeq(menuCount), deadString, getSignedNumber(differences[deadString]), getColor(differences[deadString]), linesep))
    result += ('%s%s: %s%s%s' % (
        getMenuSeq(menuCount), curedString, getSignedNumber(differences[curedString]), getColor(differences[curedString]), linesep))
    if nextMenu:
        menuCount -= 1
    result += ('%s%s' % (getPartLine(menuCount), linesep))
    return result


def showProvinceInfo(province, countryName):
    global textColor
    global displayString
    global updateString
    global confirmedString
    global deadString
    global curedString
    global provinceChangeNextMenu
    global cityChangeNextMenu

    provinceName = province.get('provinceShortName')
    provinceConfirmedCount = province.get('confirmedCount')
    provinceDeadCount = province.get('deadCount')
    provinceCuredCount = province.get('curedCount')

    displayString += ("%s%s %s: %s %s: %s %s: %s%s" % (getMenuSeq(0), provinceName, confirmedString, provinceConfirmedCount, deadString, provinceDeadCount, curedString, provinceCuredCount, linesep))

    displayString += getgetDifferencesString(countryName, provinceName, None, provinceConfirmedCount, 0, provinceDeadCount, provinceCuredCount, 1, provinceChangeNextMenu)

    updateString += (', "%s":{"%s":%s, "%s":%s, "%s":%s' % (
        provinceName, confirmedString, provinceConfirmedCount, deadString, provinceDeadCount, curedString,
        provinceCuredCount))

    comment = province.get("comment")
    if comment:
        displayString += ('%s%s | color=%s%s' % (getMenuSeq(1), comment, textColor, linesep))

    cityList = province.get('cities')
    for city in cityList:
        cityName = city.get('cityName')
        cityConfirmCount = city.get('confirmedCount')
        cityDeadCount = city.get('deadCount')
        cityCuredCount = city.get('curedCount')
        cityDataStr = "%s %s: %s %s: %s %s: %s" % (cityName, confirmedString, cityConfirmCount, deadString, cityDeadCount, curedString, cityCuredCount)
        updateString += (', "%s":{"%s":%s, "%s":%s, "%s":%s}' % (cityName, confirmedString, cityConfirmCount, deadString, cityDeadCount, curedString, cityCuredCount))
        displayString += ('%s%s | color=%s%s' % (getMenuSeq(1), cityDataStr, textColor, linesep))
        displayString += getgetDifferencesString(countryName, provinceName, cityName, cityConfirmCount, 0, cityDeadCount, cityCuredCount, 2, cityChangeNextMenu)

    updateString += '}'


def chinaProvince(countryName):
    global displayString

    if len(targetProvinceName) > 0:
        for province in jsondata:
            provinceName = province.get('provinceShortName')
            if provinceName in targetProvinceName:
                showProvinceInfo(province, countryName)
    else:
        for index in range(frontCount):
            province = jsondata[index]
            provinceName = province.get('provinceShortName')
            if provinceName not in additionProvinceName:
                showProvinceInfo(province, countryName)

    if len(jsondata) != 0:
        displayString += ('%s%s' % (getPartLine(0), linesep))
        for province in jsondata:
            provinceName = province.get('provinceShortName')
            if provinceName in additionProvinceName:
                showProvinceInfo(province, countryName)


def china():
    global textColor
    global displayString
    global jsondata
    global updateString
    global confirmedString
    global deadString
    global curedString
    global countryChangeNextMenu

    response = requests.get('https://3g.dxy.cn/newh5/view/pneumonia')
    response.encoding = 'utf-8'

    rawresult = re.search('<script id="getStatisticsService">(.*)</script>', response.text)
    provincedata = rawresult.group().split('catch')

    finalresult = provincedata[0].replace('<script id="getStatisticsService">try { window.getStatisticsService = ', '')
    finalresult = finalresult[0:-1]

    jsondata = json.loads(finalresult)

    chinaConfirmCount = jsondata.get('confirmedCount')
    chinaSuspectedCount = jsondata.get('suspectedCount')
    chinaDeadCount = jsondata.get('deadCount')
    chinaCuredCount = jsondata.get('curedCount')

    name = "中国"
    updateString += (', "%s":{"%s":%s, "%s":%s, "%s":%s, "%s":%s' % (name, confirmedString, chinaConfirmCount, suspectedString, chinaSuspectedCount, deadString, chinaDeadCount, curedString, chinaCuredCount))
    displayString += ('%s %s: %s %s:%s %s: %s %s: %s%s' % (name, confirmedString, chinaConfirmCount, suspectedString, chinaSuspectedCount, deadString, chinaDeadCount, curedString, chinaCuredCount, linesep))
    displayString += ('%s%s' % (getPartLine(0), linesep))
    displayString += ('上次更新时间: %s | color=%s%s' % (getRecordFormatTime(), textColor, linesep))
    displayString += ('%s' % getgetDifferencesString(name, None, None, chinaConfirmCount, chinaSuspectedCount, chinaDeadCount, chinaCuredCount, 0, countryChangeNextMenu))

    rawresult = re.search('<script id="getAreaStat">(.*)</script>', response.text)
    provincedata = re.search('\[.*\]', rawresult.group(1)).group(0).split('catch')

    finalresult = provincedata[0]
    finalresult = finalresult[0:-1]

    jsondata = json.loads(finalresult)
    chinaProvince(name)
    updateString += '}'


def display():
    print(displayString.strip())


def main():
    initLastData()
    readLastData()
    china()
    display()
    writeLastData()


if __name__ == "__main__":
    main()
